package algorithms.leaning.class35;

/**
 * 有序表中AVL树的实现
 * 任何一棵树的左右孩子高度不可超过1
 *
 * @author guichang
 * @date 2021/7/11
 */

@SuppressWarnings("all")
public class Code1_有序表_AVL树 {

    /**
     * 树的节点对象
     *
     * @param <K> 排序的标准，所以必要要实现Comparable接口
     * @param <V> 实际存储值，key相同会被替换，因为AVL树不可包含相同元素
     */
    public static class AVLNode<K extends Comparable<K>, V> {
        public K k;
        public V v;
        public AVLNode<K, V> l;
        public AVLNode<K, V> r;
        // 自己的树高度
        public int h;

        public AVLNode(K key, V value) {
            k = key;
            v = value;
            h = 1;
        }
    }

    /**
     * AVL树对象
     */
    public static class AVLTreeMap<K extends Comparable<K>, V> {
        // 持有根节点就可以了
        private AVLNode<K, V> root;
        // 节点个数
        private int size;

        public AVLTreeMap() {
            root = null;
            size = 0;
        }

        /**
         * 当前的cur节点右旋，返回新头部
         */
        private AVLNode<K, V> rightRotate(AVLNode<K, V> cur) {
            // 记录自己的左孩子left，left要升上去
            AVLNode<K, V> left = cur.l;
            // left的右孩子挂到cur的左孩子下
            cur.l = left.r;
            // left替代cur并且cur挂在left的右孩子下
            left.r = cur;
            // 重新计算节点高度
            cur.h = Math.max((cur.l != null ? cur.l.h : 0), (cur.r != null ? cur.r.h : 0)) + 1;
            left.h = Math.max((left.l != null ? left.l.h : 0), (left.r != null ? left.r.h : 0)) + 1;
            // 返回新头部
            return left;
        }

        /**
         * 当前的cur节点左旋。返回新头部
         */
        private AVLNode<K, V> leftRotate(AVLNode<K, V> cur) {
            // 记录自己的右孩子right，right要升上去
            AVLNode<K, V> right = cur.r;
            // right的左孩子挂到cur的右孩子下
            cur.r = right.l;
            // right替代cur并且cur挂在right的左孩子下
            right.l = cur;
            // 重新计算节点高度
            cur.h = Math.max((cur.l != null ? cur.l.h : 0), (cur.r != null ? cur.r.h : 0)) + 1;
            right.h = Math.max((right.l != null ? right.l.h : 0), (right.r != null ? right.r.h : 0)) + 1;
            // 返回新头部
            return right;
        }

        /**
         * 调整curr树以适应AVL规则，返回新头部
         */
        private AVLNode<K, V> maintain(AVLNode<K, V> cur) {
            if (cur == null) {
                return null;
            }
            // 左树高度
            int leftHeight = cur.l != null ? cur.l.h : 0;
            // 右树高度
            int rightHeight = cur.r != null ? cur.r.h : 0;
            // 高度差绝对值超过1表示要调整了
            if (Math.abs(leftHeight - rightHeight) > 1) {
                // 左树高
                if (leftHeight > rightHeight) {
                    // 左树左节点高度
                    int leftLeftHeight = cur.l != null && cur.l.l != null ? cur.l.l.h : 0;
                    // 左树右节点高度
                    int leftRightHeight = cur.l != null && cur.l.r != null ? cur.l.r.h : 0;
                    // LL型或LL和LR型同时存在结构统一按LL型调整
                    if (leftLeftHeight >= leftRightHeight) {
                        // 当前节点右旋即可
                        cur = rightRotate(cur);
                    } else {
                        // 单独时LR型的，先将左树进行一次左旋
                        cur.l = leftRotate(cur.l);
                        // 再将当前节点进行一次右旋
                        cur = rightRotate(cur);
                    }
                }
                // 右树高 操作类似
                else {
                    int rightLeftHeight = cur.r != null && cur.r.l != null ? cur.r.l.h : 0;
                    int rightRightHeight = cur.r != null && cur.r.r != null ? cur.r.r.h : 0;
                    if (rightRightHeight >= rightLeftHeight) {
                        cur = leftRotate(cur);
                    } else {
                        cur.r = rightRotate(cur.r);
                        cur = leftRotate(cur);
                    }
                }
            }
            return cur;
        }

        private AVLNode<K, V> findLastIndex(K key) {
            AVLNode<K, V> pre = root;
            AVLNode<K, V> cur = root;
            while (cur != null) {
                pre = cur;
                if (key.compareTo(cur.k) == 0) {
                    break;
                } else if (key.compareTo(cur.k) < 0) {
                    cur = cur.l;
                } else {
                    cur = cur.r;
                }
            }
            return pre;
        }

        private AVLNode<K, V> findLastNoSmallIndex(K key) {
            AVLNode<K, V> ans = null;
            AVLNode<K, V> cur = root;
            while (cur != null) {
                if (key.compareTo(cur.k) == 0) {
                    ans = cur;
                    break;
                } else if (key.compareTo(cur.k) < 0) {
                    ans = cur;
                    cur = cur.l;
                } else {
                    cur = cur.r;
                }
            }
            return ans;
        }

        private AVLNode<K, V> findLastNoBigIndex(K key) {
            AVLNode<K, V> ans = null;
            AVLNode<K, V> cur = root;
            while (cur != null) {
                if (key.compareTo(cur.k) == 0) {
                    ans = cur;
                    break;
                } else if (key.compareTo(cur.k) < 0) {
                    cur = cur.l;
                } else {
                    ans = cur;
                    cur = cur.r;
                }
            }
            return ans;
        }

        /**
         * 增加一个新节点（不重复）
         */
        private AVLNode<K, V> add(AVLNode<K, V> cur, K key, V value) {
            // 当前时空树直接创建即可
            if (cur == null) {
                return new AVLNode<K, V>(key, value);
            }
            // 当前不是空树，新加的key比当前节点小在左树加，否则去右树加
            if (key.compareTo(cur.k) < 0) {
                cur.l = add(cur.l, key, value);
            } else {
                cur.r = add(cur.r, key, value);
            }
            // 调整高度
            cur.h = Math.max(cur.l != null ? cur.l.h : 0, cur.r != null ? cur.r.h : 0) + 1;
            // 调整符合AVL树结构
            return maintain(cur);
        }

        /**
         * 在cur这棵树上，删掉key所代表的节点，返回cur这棵树的新头部
         */
        private AVLNode<K, V> delete(AVLNode<K, V> cur, K key) {
            // 与新增同理，要删除的key比当前key小去左树删，否则去右树删
            if (key.compareTo(cur.k) > 0) {
                cur.r = delete(cur.r, key);
            } else if (key.compareTo(cur.k) < 0) {
                cur.l = delete(cur.l, key);
            }
            // 找到key了
            else {
                // 是叶子节点，当前节点直接置为null
                if (cur.l == null && cur.r == null) {
                    cur = null;
                }
                // 有左无右或有右无左则将存在的节点替换
                else if (cur.l == null && cur.r != null) {
                    cur = cur.r;
                } else if (cur.l != null && cur.r == null) {
                    cur = cur.l;
                }
                // 左右孩子都有
                else {
                    // 找到右树最左节点des
                    AVLNode<K, V> des = cur.r;
                    while (des.l != null) {
                        des = des.l;
                    }
                    // 先删掉并调整好右树结构
                    cur.r = delete(cur.r, des.k);
                    // des替换cur
                    des.l = cur.l;
                    des.r = cur.r;
                    cur = des;
                }
            }
            // 调整高度
            if (cur != null) {
                cur.h = Math.max(cur.l != null ? cur.l.h : 0, cur.r != null ? cur.r.h : 0) + 1;
            }
            // 调整符合AVL结构
            return maintain(cur);
        }

        public int size() {
            return size;
        }

        public boolean containsKey(K key) {
            if (key == null) {
                return false;
            }
            AVLNode<K, V> lastNode = findLastIndex(key);
            return lastNode != null && key.compareTo(lastNode.k) == 0 ? true : false;
        }

        public void put(K key, V value) {
            if (key == null) {
                return;
            }
            AVLNode<K, V> lastNode = findLastIndex(key);
            if (lastNode != null && key.compareTo(lastNode.k) == 0) {
                lastNode.v = value;
            } else {
                size++;
                root = add(root, key, value);
            }
        }

        public void remove(K key) {
            if (key == null) {
                return;
            }
            if (containsKey(key)) {
                size--;
                root = delete(root, key);
            }
        }

        public V get(K key) {
            if (key == null) {
                return null;
            }
            AVLNode<K, V> lastNode = findLastIndex(key);
            if (lastNode != null && key.compareTo(lastNode.k) == 0) {
                return lastNode.v;
            }
            return null;
        }

        public K firstKey() {
            if (root == null) {
                return null;
            }
            AVLNode<K, V> cur = root;
            while (cur.l != null) {
                cur = cur.l;
            }
            return cur.k;
        }

        public K lastKey() {
            if (root == null) {
                return null;
            }
            AVLNode<K, V> cur = root;
            while (cur.r != null) {
                cur = cur.r;
            }
            return cur.k;
        }

        public K floorKey(K key) {
            if (key == null) {
                return null;
            }
            AVLNode<K, V> lastNoBigNode = findLastNoBigIndex(key);
            return lastNoBigNode == null ? null : lastNoBigNode.k;
        }

        public K ceilingKey(K key) {
            if (key == null) {
                return null;
            }
            AVLNode<K, V> lastNoSmallNode = findLastNoSmallIndex(key);
            return lastNoSmallNode == null ? null : lastNoSmallNode.k;
        }

    }

}